#!/usr/bin/python3
# -*- coding: utf-8 -*-

import os
import json
import tqdm
from xml.dom import minidom


def load_task_detections(categories, annotations, detection_file,id_file_map,id_categorie_map):
    """ 加载task格式的标注数据
    """
    with open(detection_file) as file:
        data = json.load(file)

    if len(data) == 0:
        return annotations

    for res in tqdm.tqdm(data):
        image_id = res["image_id"]
        image_name = id_file_map[image_id]#id to file
        if image_name not in annotations:
            print("DETECTION LOAD: not find image {} in annotations".format(image_name))
            continue

        cls_id = res["category_id"]
        cls_name = id_categorie_map[cls_id]
        if not cls_name in categories:
            print("DETECTION LOAD: can not find category {} in annos".format(cls_name))
            continue
        bBox = res["bbox"]
        bBox[2] = bBox[0]+bBox[2]
        bBox[3] = bBox[1]+bBox[3]#"bbox": [xmin, ymin, xmax, ymax],

        bboxes = {"bbox": bBox, "category": cls_name, "score": res["score"]}
        annotations[image_name]["dt"].append(bboxes)

    return annotations



def load_detections(categories, annotations, detection_file,id_file_map,id_categorie_map):
    """
    输出：
    [
        {
            "image_name": {
                "gt": [anno1, anno2, ...],
                "dt": [det1, det2, ...]
            },
            ...
        }
    ]

    det格式：
    {
        "bbox": [xmin, ymin, xmax, ymax],
        "category": cls_name,
        "score": s
    }
    """
    print("loading detections {}".format(detection_file))

    annotations = load_task_detections(categories, annotations, detection_file,id_file_map,id_categorie_map)

    # 在匹配的时候通过阈值判断检测数量
    # for image_name in annotations:
    #     for det in annotations[image_name]["dt"]:
    #        categories[det["category"]]["dtn"] += 1


    det_images_cnt = 0
    det_objects_cnt = 0
    for name in annotations:
        if len(annotations[name]["dt"]) > 0:
            det_images_cnt += 1
            det_objects_cnt += len(annotations[name]["dt"])
    print("load {} detections from {} images".format(det_objects_cnt, det_images_cnt))

    return categories, annotations